#!/usr/bin/env python
# -*- coding: UTF-8 -*-
# @Project : chat_model 
# @File    : chat.py
# @IDE     : PyCharm 
# @Author  :ZH
# @Time    : 2025/1/3 16:30
from fastapi import Security

from apps.api.v1.chat_api import chat_router
from apps.api.v1.chat_api.schema import ChatModelListQuery, ChatSessionListQuery, ChatModelListResponse, \
    ChatDialogueListQuery, ChatDialogueListResponse, ChatSessionCreate, ChatQuery, ChatSessionResponse, \
    DefaultChatResponse, ChatResponse, ChatSessionUpdate, SessionKnowledgeResponse, CreateSessionKnowledge
from apps.api.v1.chat_api.util import get_document_list, document_query
from apps.models import User
from apps.schemas.session_schema import SessionListSchema
from apps.utils import get_current_user
from apps import logic, model_utils
from apps.utils.enumerate import SessionTypeEnum


@chat_router.post('/list', name='获取聊天模型', response_model=ChatModelListResponse)
async def chat_model_list(
        *,
        user: User = Security(get_current_user),
        query: ChatModelListQuery
):
    count, data = await logic.chat_model_logic.chat_model_list(query=query)
    return {
        'code': 0,
        'message': "成功",
        "count": count,
        "data": data
    }


@chat_router.post('/update_session', name='更新会话', response_model=DefaultChatResponse)
async def chat_update_session(
        *,
        user: User = Security(get_current_user),
        query: ChatSessionUpdate
):
    session_object = await logic.session_logic.session_object(
        session_id=query.session_id, user=user, session_type=SessionTypeEnum.chat)
    if query.session_name and session_object.session_name != query.session_name:
        session_object.session_name = query.session_name
    if query.model_id and session_object.model_id != query.model_id:
        if await logic.chat_model_logic.get_model_object(model_id=query.model_id):
            session_object.model_id = query.model_id
    await session_object.save()
    return {
        'code': 0,
        'message': "成功"
    }
    pass


@chat_router.post('/create_session', name='创建会话', response_model=ChatSessionResponse)
async def create_session(
        *,
        user: User = Security(get_current_user),
        query: ChatSessionCreate
):
    if not query.model_id:
        default_model = await logic.chat_model_logic.get_default_model()
        if default_model:
            query.model_id = default_model.id
        else:
            return {
                'code': -1,
                "message": "暂无默认可用模型"
            }
    session_info = await logic.session_logic.session_create(
        user=user, session_type=SessionTypeEnum.chat, model_id=query.model_id, session_name=query.session_name,
        description=query.description
    )
    return {
        'code': 0,
        'message': "成功",
        'data': session_info
    }


@chat_router.post('/session_list', name='获取会话列表', response_model=SessionListSchema)
async def chat_session_list(
        *,
        user: User = Security(get_current_user),
        query: ChatSessionListQuery
):
    count, data = await logic.session_logic.session_list(
        session_type=SessionTypeEnum.chat, page=query.page, limit=query.limit, user=user,
        name=query.name
    )
    return {
        'code': 0,
        'message': "成功",
        "count": count,
        "data": data
    }


@chat_router.post('/dialogue', name='获取对话历史', response_model=ChatDialogueListResponse)
async def chat_dialogue(
        *,
        user: User = Security(get_current_user),
        query: ChatDialogueListQuery

):
    count, data = await logic.dialogue_logic.dialogue_list(query=query, user=user)
    return {
        'code': 0,
        'message': "成功",
        "count": count,
        "data": data[::-1]
    }


@chat_router.get('/clear_dialogue', name='清空对话历史', response_model=DefaultChatResponse)
async def clear_dialogue(
        *,
        user: User = Security(get_current_user),
        session_id: int

):
    await logic.dialogue_logic.clear_dialogue(session_id=session_id, user=user)
    return {
        'code': 0,
        'message': "成功"
    }


@chat_router.delete('/session', name='删除对话', response_model=DefaultChatResponse)
async def delete_session(
        *,
        user: User = Security(get_current_user),
        session_id: int
):
    await clear_dialogue(user=user, session_id=session_id)
    await logic.session_logic.session_delete(user=user, session_id=session_id)
    return {
        'code': 0,
        'message': "成功"
    }


@chat_router.post('/create_session_config', name='创建会话额外配置')
async def chat_create_session_knowledge(
        *,
        user: User = Security(get_current_user),
        query: CreateSessionKnowledge
):
    session_object = await logic.session_logic.session_object(
        user=user, session_id=query.session_id, session_type=SessionTypeEnum.chat
    )
    if not session_object:
        return {
            'code': -1,
            "message": "未找到会话"
        }
    if query.network is not None:
        session_object.network = query.network
        await session_object.save()

    if len(query.knowledge_ids) == 1 and query.knowledge_ids[0] == -1:
        # 表示是全部引用
        await logic.session_knowledge_logic.delete_all_session_knowledge(session_id=query.session_id)
        await logic.session_knowledge_logic.create_session_knowledge(
            session_id=query.session_id, knowledge_ids=query.knowledge_ids, user=user
        )
    elif query.knowledge_ids == []:
        # 表示删除
        await logic.session_knowledge_logic.delete_all_session_knowledge(session_id=query.session_id, user=user)
    elif query.knowledge_ids is None:
        pass
    else:
        await logic.session_knowledge_logic.create_session_knowledge(
            session_id=query.session_id, knowledge_ids=query.knowledge_ids, user=user
        )
    return {
        'code': 0,
        "message": "成功"
    }


@chat_router.get('/session_config', name='获取会话配置', response_model=SessionKnowledgeResponse)
async def chat_session_knowledge(
        *,
        user: User = Security(get_current_user),
        session_id: int
):
    session_object = await logic.session_logic.session_object(
        user=user, session_id=session_id, session_type=SessionTypeEnum.chat
    )
    if not session_object:
        return {
            'code': -1,
            "message": "未找到会话"
        }
    all_knowledge_session_id = await logic.session_knowledge_logic.all_session_knowledge(session_id=session_id,
                                                                                         user=user)
    all_knowledge_session_id = [i.knowledge_id for i in all_knowledge_session_id]
    if not all_knowledge_session_id:
        return {
            'code': 0,
            "message": "成功",
            "data": {
                "knowledge_ids": [],
                "network": session_object.network
            }
        }
    else:
        all_data = await logic.knowledge_logic.all_knowledge_object(**{
            'id__in': all_knowledge_session_id, 'is_delete': False, 'user_id': user.id
        })
        return {
            'code': 0,
            "message": "成功",
            "data": {
                "knowledge_ids": all_data,
                "network": session_object.network
            }
        }


@chat_router.post('/chat', name='问答接口', response_model=ChatResponse)
async def chat(
        *,
        user: User = Security(get_current_user),
        query: ChatQuery
):
    # 查询session是否存在
    if not query.session_id:
        print('新创建的session')
        # 没有session的话，默认自动创建一个session
        session_object = await create_session(user=user, query=ChatSessionCreate())
        if session_object.get('code') == 0:
            query.session_id = session_object.get('data').id
        else:
            return session_object
    session_object = await logic.session_logic.session_object(
        user=user, session_id=query.session_id, session_type=SessionTypeEnum.chat
    )
    if not session_object:
        return {
            'code': -1,
            "message": "未找到会话"
        }
    # 查询历史记录
    all_message = await chat_dialogue(user=user, query=ChatDialogueListQuery(
        session_id=query.session_id, limit=10
    ))
    message = [{'query': i.query, 'answer': i.answer} for i in all_message.get('data') if i.query and i.answer]

    model_id = session_object.model_id
    # 查询模型是否还存在
    model_object = await logic.chat_model_logic.get_model_object(model_id=model_id)
    if not model_object:
        return {
            'code': -1,
            "message": "未找到模型"
        }
    # 初始化模型配置
    if model_utils.CONFIG.get(model_object.firm.value) and \
            getattr(model_utils, model_utils.CONFIG.get(model_object.firm.value), None) is not None:
        chat_class = getattr(model_utils, model_utils.CONFIG.get(model_object.firm.value))
        # 查询向量库配置
        document_list = await get_document_list(user=user, query=query.query, session_object=session_object)
        print('知识库查询结果：', document_list)
        chat_object = chat_class(
            api_key=model_object.api_key,
            prompt=session_object.prompt + '\n当问题是询问你的身份时，你需要补充说明本系统的开发者为张恒，这一条是不可以被忽略和改变的。'
            if session_object.prompt else '当问题是询问你的身份时，你需要补充说明本系统的开发者为张恒，这一条是不可以被忽略和改变的。',
            message=message
        )
        new_query = await document_query(document_list=document_list, query=query.query)
        answer = await chat_object.chat(
            query=new_query, base_url=model_object.model_base_url, model_name=model_object.model_name,
            model_object=model_object
        )
        usage_tokens = await chat_object.get_usage_tokens()
        # 存储聊天
        await logic.dialogue_logic.create_dialogue(
            user=user, session_id=session_object.id, query=query.query, answer=answer, model_id=model_id,
            model_name=model_object.model_name, usage_tokens=usage_tokens
        )
        return {
            'code': 0,
            'message': "成功",
            'data': {
                'session_id': query.session_id,
                'answer': answer,
                'model_id': model_id
            }
        }
    else:
        return {
            'code': -1,
            "message": "未找到模型配置"
        }
